package org.example.springboot.common.aspect;

import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.example.springboot.common.annotation.RateLimit;
import org.example.springboot.common.exception.ServerException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.time.Instant;
import java.util.Collections;
import java.util.concurrent.TimeUnit;

@Aspect
@Component
public class RateLimitAspect {

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    @Resource
    private HttpServletRequest request;

    @Pointcut("@annotation(org.example.springboot.common.annotation.RateLimit)")
    public void rateLimit() {
    }

    @Around("rateLimit()")
    public Object ponitCut(ProceedingJoinPoint point) throws Throwable {
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();

        RateLimit rateLimit = AnnotationUtils.findAnnotation(method, RateLimit.class);
        if (rateLimit != null) {
            String redisKey = method.getDeclaringClass().getName() + "." + method.getName();
            // CacheConstants为常量类，主要用来定义redis的key前缀名
            // RequestUtil.getIp()为获取ip的方法，可根据需求自定义
            redisKey += "smart" + "-" + getClientIP(request);

            long max = rateLimit.max(); // api最大请求次数
            long timeout = rateLimit.timeout(); // api请求次数的指定时间(秒)
            long ttl = TimeUnit.SECONDS.toMillis(timeout); // 转换成毫秒
            long nowTime = Instant.now().toEpochMilli(); // 当前时间毫秒
            long expire = nowTime - ttl; // 计算过期时间
            DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
            redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("redis/RateLimit.lua")));
            redisScript.setResultType(Long.class);
            // 执行lua脚本，返回已请求api的次数，0表示超出请求次数
            long execCount = stringRedisTemplate.execute(redisScript, Collections.singletonList(redisKey),
                    nowTime + "", ttl + "", expire + "", max + "");
            if (execCount == 0) {
                throw new ServerException("请求次数已达上限，请稍后再试");
            }
        }

        return point.proceed();
    }


    public String getClientIP(HttpServletRequest request) {
        String ip = request.getHeader("X-Forwarded-For");
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }
        return ip;
    }

}